import matplotlib.pyplot as plt
import numpy as np
font = {'family': 'Arial', 'size' : 16}
plt.rc('font', **font)

plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.it'] = 'Arial:italic'
plt.rcParams['mathtext.rm'] = 'Arial'
plt.rcParams['pdf.fonttype'] = 42

# Fixing random state for reproducibility
np.random.seed(19680801)

# fig, axs = plt.subplots(2, 2)
# cmaps = ['RdBu_r', 'viridis']
# for col in range(2):
#     for row in range(2):
#         ax = axs[row, col]
#         pcm = ax.pcolormesh(np.random.random((20, 20)) * (col + 1),
#                             cmap=cmaps[col])
#         fig.colorbar(pcm, ax=ax)

# linear noise
data0 = np.zeros((4, 5))
print(data0)
data0[0] = [0.40307, 0.45477, 0.58309, 0.64967, 0.76351] # no noise
data0[1] = [0.55491, 0.61395, 0.67604, 0.66686, 0.77571] # low noise
data0[2] = [0.85933, 0.9456, 0.80721, 0.83281, 0.78227] # medium noise
data0[3] = [1.14253, 1.31645, 1.01337, 0.8222, 0.76578] # high noise

print(data0)
# normal noise
data1 = np.zeros((4, 5))
data1[0] = [0.40307, 0.45477, 0.58309, 0.64967, 0.76351] # no noise
data1[1] = [0.60644, 0.69187, 0.66626, 0.69316, 0.77087] # low noise
data1[2] = [0.80749, 0.96945, 0.91626, 0.76517, 0.76853] # medium noise
data1[3] = [0.92668, 1.03641, 1.00043, 0.80959, 0.76479] # high noise

fig, axs = plt.subplots(1, 2, constrained_layout=True, sharey=True, figsize=(10, 5))
for i in range(4):
    for j in range(5):
        axs[0].text(j + 0.5, i + 0.5, "%.2f" % data0[i][j], ha='center', color=["white","black"][data0[i][j]>0 and data0[i][j]<0.9],
                    va='center')
        # if i == 0:
        #     continue
        # axs[0].text(j + 0.5, i, "%.0f" % ((data0[i][j]/data0[0][j] - 1) *100), ha='center', color=["white","black"][data0[i][j]>0 and data0[i][j]<0.9],
        #             va='center')


for i in range(4):
    for j in range(5):
        axs[1].text(j + 0.5, i + 0.5, "%.2f" % data1[i][j], ha='center', color=["white","black"][data1[i][j]>0 and data1[i][j]<0.9],
                    va='center')
        # if i == 0:
        #     continue
        # axs[1].text(j + 0.5, i, "%.0f" % ((data1[i][j]/data1[0][j] - 1) *100), ha='center', color=["white","black"][data0[i][j]>0 and data0[i][j]<0.9],
        #             va='center')




axs[0].set_ylabel("Noise level")
axs[0].annotate("(a)", xy=(0, 0), xytext=(-0.4, 3.8))
axs[1].annotate("(b)", xy=(0, 0), xytext=(-0.4, 3.8))
for ax in axs.flat:
    ax.set_xlabel("Dataset size")
    ax.set_xticks([0.5, 1.5, 2.5, 3.5, 4.5], labels=['52348\nABCD', '30000\nA', '10000\nB', '5000\nC', '2348\nD'])
    ax.set_yticks([0.5, 1.5, 2.5, 3.5], labels=['None', 'Low', 'Med', 'High'])


pcm = axs[0].pcolormesh(data0, vmin=0.4, vmax=1.3, cmap='Greys')
pcm = axs[1].pcolormesh(data1, vmin=0.4, vmax=1.3, cmap='Greys')

cb = fig.colorbar(pcm, ax=[axs[0], axs[1]], shrink=0.6, location='bottom')
cb.set_ticks([0.4, 0.6, 0.8, 1.0, 1.2], labels=['0.4', '0.6', '0.8', '1.0', '1.2 (eV)'])
# fig.colorbar(pcm, ax=[axs[0, 2]], location='bottom')
# fig.colorbar(pcm, ax=axs[1:, :], location='right', shrink=0.6)
# fig.colorbar(pcm, ax=[axs[2, 1]], location='left')

plt.show()